{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# 如何设计&开发新的明密文混合算法"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "> This tutorial is only available in Chinese.\n",
    "\n",
    "> 以下代码仅作为示例，请勿在生产环境直接使用。"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## 集群部署模式"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "vscode": {
     "languageId": "plaintext"
    }
   },
   "source": [
    "隐语有多种集群部署模式，包括[仿真模式](../getting_started/deployment.md#simulation)、[生产模式](../getting_started/deployment.md#production)和[调试模式](../getting_started/debug_mode.md)。不同的集群部署模式对代码编写几乎没有影响，只有集群初始化的代码有所不同。因此本教程中的示例都是以仿真模式编写和运行的。\n",
    "\n",
    "下面首先以单机仿真模式启动一个4节点的集群："
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 1,
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "2023-11-06 20:03:44,807\tINFO worker.py:1538 -- Started a local Ray instance.\n"
     ]
    }
   ],
   "source": [
    "import secretflow as sf\n",
    "\n",
    "# In case you have a running SecretFlow runtime already.\n",
    "sf.shutdown()\n",
    "\n",
    "sf.init(['alice', 'bob', 'carol', 'dave'], address='local')"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Device 概念 "
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "在开始前建议先阅读 SecretFlow 的[架构设计](../developer/design/architecture.md)。\n",
    "\n",
    "在集群初始化时，我们声明了4个参与方：alice、bob、carol、dave，它们对应4个“物理设备”。在 SecretFlow 中物理设备用来直观的表示参与方和组网信息，用于实际执行计算和网络通信。\n",
    "\n",
    "隐语在物理设备上构建了一套逻辑设备：PYU、SPU、HEU、TEEU。一个逻辑设备由一个或多个物理设备组成，同时一个物理设备也可以用于多个逻辑设备。在编程意义上，逻辑设备负责管理“数据”和执行“计算”，在隐语上编程可以简单理解为 3 个部分：\n",
    "1. 像明文编程一样编写逻辑设备的计算过程；\n",
    "2. 将数据放到逻辑设备上，或者在逻辑设备之间流转；\n",
    "3. 控制逻辑设备，使用该逻辑设备上的“数据”执行一段计算。\n",
    "\n",
    "可以说隐语上的编程就是面向逻辑设备的编程。"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "(PYURuntime(alice),\n",
       " PYURuntime(bob),\n",
       " '<secretflow.device.device.pyu.PYU object at 0x7f4a9b8c0e80>')"
      ]
     },
     "execution_count": 2,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "# 声明两个 PYU 逻辑设备，alice 和 bob，分别代表在对应的物理设备本地计算\n",
    "alice, bob = sf.PYU('alice'), sf.PYU('bob')\n",
    "\n",
    "alice, bob, object.__repr__(alice)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "<secretflow.device.device.spu.SPU at 0x7f4a9b8add90>"
      ]
     },
     "execution_count": 3,
     "metadata": {},
     "output_type": "execute_result"
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "\u001b[2m\u001b[36m(SPURuntime pid=3915774)\u001b[0m 2023-11-06 20:05:54.079 [info] [default_brpc_retry_policy.cc:DoRetry:52] socket error, sleep=1000000us and retry\n",
      "\u001b[2m\u001b[36m(SPURuntime pid=3915774)\u001b[0m 2023-11-06 20:05:55.079 [info] [default_brpc_retry_policy.cc:LogHttpDetail:29] cntl ErrorCode '112', http status code '200', response header '', error msg '[E111]Fail to connect Socket{id=0 addr=127.0.0.1:54955} (0x0x335c640): Connection refused [R1][E112]Not connected to 127.0.0.1:54955 yet, server_id=0'\n",
      "\u001b[2m\u001b[36m(SPURuntime pid=3915774)\u001b[0m 2023-11-06 20:05:55.079 [info] [default_brpc_retry_policy.cc:DoRetry:75] aggressive retry, sleep=1000000us and retry\n",
      "\u001b[2m\u001b[36m(SPURuntime pid=3915774)\u001b[0m 2023-11-06 20:05:56.079 [info] [default_brpc_retry_policy.cc:LogHttpDetail:29] cntl ErrorCode '112', http status code '200', response header '', error msg '[E111]Fail to connect Socket{id=0 addr=127.0.0.1:54955} (0x0x335c640): Connection refused [R1][E112]Not connected to 127.0.0.1:54955 yet, server_id=0 [R2][E112]Not connected to 127.0.0.1:54955 yet, server_id=0'\n",
      "\u001b[2m\u001b[36m(SPURuntime pid=3915774)\u001b[0m 2023-11-06 20:05:56.080 [info] [default_brpc_retry_policy.cc:DoRetry:75] aggressive retry, sleep=1000000us and retry\n",
      "\u001b[2m\u001b[36m(SPURuntime pid=3915775)\u001b[0m 2023-11-06 20:05:56.135 [info] [default_brpc_retry_policy.cc:DoRetry:69] not retry for reached rcp timeout, ErrorCode '1008', error msg '[E1008]Reached timeout=2000ms @127.0.0.1:41437'\n"
     ]
    }
   ],
   "source": [
    "# 使用两个物理设备组成一个 SPU 逻辑设备\n",
    "spu = sf.SPU(sf.utils.testing.cluster_def(['carol', 'dave']))\n",
    "\n",
    "spu"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## 在 PYU 明文设备上编程"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "明文设备和物理设备对应，用于编写和运行在单方明文执行的代码。\n",
    "下面以一个多方求均值的例子演示一下如何在 PYU 设备上编写代码。\n",
    "\n",
    "任务：alice、bob、carol 三方各自有一些整型数值， 目标是计算出三方所有数字的均值，dave 是一个可用的辅助计算方。"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### 声明 PYU 设备 "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "metadata": {},
   "outputs": [],
   "source": [
    "alice = sf.PYU('alice')\n",
    "bob = sf.PYU('bob')\n",
    "carol = sf.PYU('carol')\n",
    "dave = sf.PYU('dave')"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### 准备数据"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "metadata": {},
   "outputs": [],
   "source": [
    "import numpy as np\n",
    "\n",
    "\n",
    "# 定义数据生成函数\n",
    "def gen_data(fake_seed, file_name):\n",
    "    import pathlib\n",
    "\n",
    "    pathlib.Path(file_name).parent.mkdir(parents=True, exist_ok=True)\n",
    "\n",
    "    data = np.arange(fake_seed, 2 * fake_seed - 1, 1)\n",
    "    np.savetxt(file_name, data, fmt='%d')"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "三方分别执行 `gen_data`，并把生成的数据保存的各自的本地，由于目前是单机仿真模式，所以实际上都会保存在本机上。"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "<secretflow.device.device.pyu.PYUObject at 0x7f4bfef4c880>"
      ]
     },
     "execution_count": 6,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "alice(gen_data)(6, './data/a.csv')\n",
    "bob(gen_data)(10, './data/b.csv')\n",
    "carol(gen_data)(11, './data/c.csv')"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "可以看到上面输出了一个 `PYUObject`，他是最后一行（也就是 carol）执行的返回值。不管实际执行的函数有没有返回值，交给 `PYU` 执行之后都会返回 `PYUObject`。这里返回的 `PYUObject` 实际值是 `None`，它虽然没有内容，但是可以用于 `sf.wait` 来等待函数执行完成，有时候是很有用的。\n",
    "\n",
    "我们来看一下生成的3份数据，后面我们将使用他们来完成计算均值的任务："
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "['6', '7', '8', '9', '10']"
      ]
     },
     "execution_count": 7,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "%%!\n",
    "cat ./data/a.csv"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "['10', '11', '12', '13', '14', '15', '16', '17', '18']"
      ]
     },
     "execution_count": 8,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "%%!\n",
    "cat ./data/b.csv"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 9,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "['11', '12', '13', '14', '15', '16', '17', '18', '19', '20']"
      ]
     },
     "execution_count": 9,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "%%!\n",
    "cat ./data/c.csv"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### 加载数据"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 10,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "(<secretflow.device.device.pyu.PYUObject at 0x7f4a9b8c08b0>,\n",
       " <secretflow.device.device.pyu.PYUObject at 0x7f4bfefa3ac0>,\n",
       " <secretflow.device.device.pyu.PYUObject at 0x7f4bfef4b2b0>)"
      ]
     },
     "execution_count": 10,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "def load_data(file_name):\n",
    "    data = np.loadtxt(file_name, dtype=np.int32)\n",
    "    return data\n",
    "\n",
    "\n",
    "alice_data = alice(load_data)('./data/a.csv')\n",
    "bob_data = bob(load_data)('./data/b.csv')\n",
    "carol_data = carol(load_data)('./data/c.csv')\n",
    "alice_data, bob_data, carol_data"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "上面代码已经把数据加载到了各方内存中，我们尝试着看一下数据`alice_data, bob_data, carol_data`，可以看到他们是一个`secretflow.device.device.pyu.PYUObject`的实例。\n",
    "\n",
    "`PYUObject` 代表了对一个 PYU 设备上的数据的引用，它本身不包含数据，我们尝试看一下`alice_data`的内容，可以看到它包含一个引用 `ObjectRef` 和他所在的 PYU `PYURuntime(alice)`(不同的组网模式会有一些不同)。"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 11,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "{'data': ObjectRef(f4402ec78d3a2607ffffffffffffffffffffffff0100000001000000),\n",
      " 'device': PYURuntime(alice)}\n"
     ]
    }
   ],
   "source": [
    "from pprint import pprint\n",
    "\n",
    "pprint(vars(alice_data))"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "这里需要注意两点：\n",
    "1. 逻辑上，一个 `PYUObject` 只能被它所在的 `PYU` 访问和处理，其他逻辑设备无法访问；\n",
    "2. 物理上，一个 `PYUObject` 只存在于它所在的 `PYU` 对应的物理设备的内存中，其他物理设备无法访问（生产模式下）。\n",
    "\n",
    "在开发过程中，我们可以用 `sf.reveal`（也就是把数据公开出来）的方式拿到真实的数据，然后打印出来。\n",
    "\n",
    "**特别注意：在生产代码中 `sf.reveal` 几乎不应该被使用，只有在算法中明确需要公开的数据才应该使用它**"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 12,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "[ 6  7  8  9 10]\n"
     ]
    }
   ],
   "source": [
    "print(sf.reveal(alice_data))"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "但在线上运行的时候我们不可能直接公开数据，但是又有打印中间数据的需求。我们仍然是可以做到这一点的，但是只能在数据所在的 `PYU` 上打印。\n",
    "\n",
    "执行下面的代码，可以看到内容是在一个其他进程中打印出来的，在仿真模式中，这个进程代表 alice 执行了`print_data`函数；如果是在生产模式，那么只有真正的 alice 方可以把内容打印出来。"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 13,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "\u001b[2m\u001b[36m(_run pid=3908759)\u001b[0m alice_data: [ 6  7  8  9 10]\n"
     ]
    }
   ],
   "source": [
    "def print_data(msg, data):\n",
    "    print(f'{msg}: {data}')\n",
    "\n",
    "\n",
    "_ = alice(print_data)('alice_data', alice_data)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### 计算均值"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "#### 设计算法"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "现在我们的三方各自有一些数据，需要算一个整体的均值，为了一定程度上保护原始数据，我们可以每方各自求和，并把求和结果和数据量发送给 dave，然后 dave 计算均值：$$avg=\\frac{sum(a)+sum(b)+sum(c)}{n_a+n_b+n_c}$$"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "#### 算法实现"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "首先三方各自求和并统计数据量，我们来分别展示一下基于 PYU 编程的几种模式。\n",
    "\n",
    "方法一：定义并执行 Task\n",
    "\n",
    "Task 是一个 PYU 的执行单元，他表示在 PYU 上执行一个 python 函数。语法如下所示，需要执行的函数可以是自定义函数，也可以是库函数或者内置函数，函数的参数放在第二个括号中，Task 返回的结果是一个 PYUObject。\n",
    "\n",
    "需要注意，下面的 `Task 1` 和 `Task 2` 两行只是提交了两个 Task，并没有等待执行完成，Task 的计算是异步完成的。"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 14,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "<secretflow.device.device.pyu.PYUObject object at 0x7f4bfef4ba00> <secretflow.device.device.pyu.PYUObject object at 0x7f4bfefc9280>\n"
     ]
    }
   ],
   "source": [
    "def sum(data):\n",
    "    return np.sum(data)\n",
    "\n",
    "\n",
    "alice_sum = alice(sum)(alice_data)  # Task 1\n",
    "alice_shape = alice(np.shape)(alice_data)  # Task 2\n",
    "\n",
    "# 看一下 Task 的返回值\n",
    "print(alice_sum, alice_shape)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "我们用 reveal 看一下结果："
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 15,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "alice sum=40, shape=(5,)\n"
     ]
    }
   ],
   "source": [
    "# wait 可以等待 PYUObject 对应的 Task 执行完成\n",
    "sf.wait([alice_sum, alice_shape])\n",
    "\n",
    "# reveal 也有等待完成的效果，但同时会把数据公开出来\n",
    "print(f'alice sum={sf.reveal(alice_sum)}, shape={sf.reveal(alice_shape)}')"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "方法二：Task 多返回值，函数可以有多返回值，可以用指定 num_returns 的方式让 Task 返回多个 PYUObject。如果不指定 num_returns，那么框架会尝试自动推导应该返回多少个 PYUObject，如果推导失败所有返回值会被封装成一个 PYUObject。除非函数特别简单，否则我们还是建议明确指定 num_returns。"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 16,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "<secretflow.device.device.pyu.PYUObject object at 0x7f4bfef48790> <secretflow.device.device.pyu.PYUObject object at 0x7f4c2453be20>\n"
     ]
    }
   ],
   "source": [
    "def agg(data):\n",
    "    return np.sum(data), np.shape(data)\n",
    "\n",
    "\n",
    "bob_sum, bob_shape = bob(agg, num_returns=2)(bob_data)\n",
    "\n",
    "print(bob_sum, bob_shape)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 17,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "bob sum=126, shape=(9,)\n"
     ]
    }
   ],
   "source": [
    "print(f'bob sum={sf.reveal(bob_sum)}, shape={sf.reveal(bob_shape)}')"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "方法三：使用 Actor 管理数据和计算。\n",
    "\n",
    "Actor 可以认为是在特定 device 上的一个类的实例，除了 `__init__` 以外，其他成员函数调用都会被封装成一个 Task，意味着成员函数的返回值会变成 PYUObject，参数也可以接受 PYUObject 或者普通变量。\n",
    "\n",
    "值得注意的是：\n",
    "1. Actor 会常驻在一个特定进程中，因此 Actor 内部可以保存状态；\n",
    "2. Actor 成员函数的参数和返回值必须是可序列化的（普通 Task 也一样），但是其内部状态没有这个要求，可以是任意数据；\n",
    "3. Actor 的 `__init__` 比较特殊，他会给原始类的 `__init__` 增加一个 `device` 参数，用来表示这个 Actor 实例在哪个 PYU 设备上;\n",
    "4. Actor 的 `__init__` 不能接收 PYUObject 作为参数，只能接收基本参数（全局共享的参数），但是 `__init__` 函数内可以正常访问 Actor 的内部状态。"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 18,
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "INFO:root:Create proxy actor <class '__main__.Stat'> with party carol.\n"
     ]
    }
   ],
   "source": [
    "from secretflow import proxy, PYUObject\n",
    "\n",
    "\n",
    "@proxy(PYUObject)\n",
    "class Stat(object):\n",
    "    def __init__(self, file_name):\n",
    "        self.file_name = file_name\n",
    "\n",
    "    def load_data(self):\n",
    "        self.data = np.loadtxt(self.file_name, dtype=np.int32)\n",
    "\n",
    "    def set_data(self, data):\n",
    "        self.data = data\n",
    "\n",
    "    def sum(self):\n",
    "        return np.sum(self.data)\n",
    "\n",
    "    def shape(self):\n",
    "        return np.shape(self.data)\n",
    "\n",
    "\n",
    "carol_actor = Stat('./data/c.csv', device=carol)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "上面我们定义了一个叫做 `Stat` 的 Actor 类，并且初始化了一个实例 `carol_actor`，我们初始化时指定了 `device=carol`，所以这个实例会在 `carol` 上运行。\n",
    "\n",
    "接下来我们就可以像普通对象一样使用 `carol_actor` 了，但是要记住，他们返回值都是 `PYUObject`："
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 19,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "carol sum=155, shape=(10,)\n"
     ]
    }
   ],
   "source": [
    "carol_actor.load_data()\n",
    "# 或者直接传入已读取的数据：\n",
    "# carol_actor.set_data(carol_data)\n",
    "carol_sum = carol_actor.sum()\n",
    "carol_shape = carol_actor.shape()\n",
    "\n",
    "print(f'carol sum={sf.reveal(carol_sum)}, shape={sf.reveal(carol_shape)}')"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "上面三方各自的计算已经完成，接下来需要把这些部分的结果汇总到 dave 上，然后在 dave 上计算出最终结果。\n",
    "\n",
    "接下来我们把数据传递到 dave 上："
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 20,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "(<secretflow.device.device.pyu.PYUObject at 0x7f4bfef4c9a0>,\n",
       " <secretflow.device.device.pyu.PYUObject at 0x7f4bfef4cdf0>)"
      ]
     },
     "execution_count": 20,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "alice_sum_d = alice_sum.to(dave)\n",
    "alice_shape_d = alice_shape.to(dave)\n",
    "\n",
    "bob_sum_d = bob_sum.to(dave)\n",
    "bob_shape_d = bob_shape.to(dave)\n",
    "\n",
    "carol_sum_d = carol_sum.to(dave)\n",
    "carol_shape_d = carol_shape.to(dave)\n",
    "\n",
    "alice_sum_d, alice_shape_d"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "然后计算均值，下面展示了另外一个小技巧：作为 Task 的参数时，PYUObject 可以放到 list 或 dict 等结构中，并且不限制层次（比如可以是 dict 中的 list 的元素）。"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 21,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "avg=13.375, type=<class 'numpy.float64'>\n"
     ]
    }
   ],
   "source": [
    "def calc_avg(sums, shapes):\n",
    "    total_sum = np.sum(sums)\n",
    "    total_shape = np.sum(shapes)\n",
    "    return total_sum / total_shape\n",
    "\n",
    "\n",
    "dave_sums = [alice_sum_d, bob_sum_d, carol_sum_d]\n",
    "dave_shapes = [alice_shape_d, bob_shape_d, carol_shape_d]\n",
    "\n",
    "dave_avg = dave(calc_avg)(dave_sums, dave_shapes)\n",
    "\n",
    "public_avg = sf.reveal(dave_avg)\n",
    "print(f'avg={sf.reveal(public_avg)}, type={type(public_avg)}')"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "最后本地计算一下，可以看到计算结果是一样的。"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 22,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "avg=[13.375], type=<class 'numpy.ndarray'>\n"
     ]
    }
   ],
   "source": [
    "data_a = np.loadtxt('./data/a.csv', dtype=np.int32)\n",
    "data_b = np.loadtxt('./data/b.csv', dtype=np.int32)\n",
    "data_c = np.loadtxt('./data/c.csv', dtype=np.int32)\n",
    "data = np.concatenate([data_a, data_b, data_c])\n",
    "\n",
    "avg = np.sum(data) / np.shape(data)\n",
    "print(f'avg={avg}, type={type(avg)}')"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "上面的本地计算中，故意使用了 `np.sum(data)/np.shape(data)` 计算均值，可以看到本地计算结果和基于 PYU 的计算结果并不完全一致：他们的类型不同。\n",
    "\n",
    "原因在于 python 的特点和 num_returns 参数：在 python 中一个函数有多个返回值时，实际上返回的是一个 tuple, 因此下面两个函数是完全等价的。同样对于 num_returns 参数来说，下面两个函数都可以设置为 `num_returns=2`。"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 23,
   "metadata": {},
   "outputs": [],
   "source": [
    "def multi_ret():\n",
    "    return 1, 2\n",
    "\n",
    "\n",
    "def multi_ret2():\n",
    "    return (1, 2)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "当返回值真的有多个时一切都很好，但是如果返回值只有一个，并且返回值是 array-like 的（tuple、list 等），且返回的 array 中只有一个值时（如下面的情况），PYU 就无法区分应该返回 array 中的元素还是 array 本身了，这种情况下，PYU 都会返回 array 中的唯一的那个元素值（也就是会解包），虽然这不一定是编程者的意图。"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 24,
   "metadata": {},
   "outputs": [],
   "source": [
    "def single_ret():\n",
    "    return 1\n",
    "\n",
    "\n",
    "def single_ret2():\n",
    "    return (1,)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "**这是一个需要被着重注意的细节：如果一个函数的返回值是一个不确定长度的 list，当返回值的元素个数为 1 的时候，它会返回了 list 内的元素本身。不做特殊处理的情况下容易出现意想不到的 bug**\n",
    "\n",
    "我们再来看一下计算过程："
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 25,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "avg=13.375, type=<class 'numpy.float64'>\n"
     ]
    }
   ],
   "source": [
    "def calc_avg(sums, shapes):\n",
    "    total_sum = np.sum(sums)\n",
    "    total_shape = np.sum(shapes)\n",
    "    # 此处返回了一个长度为 1 的 array-like 数据\n",
    "    return total_sum / total_shape  # [x]\n",
    "\n",
    "\n",
    "dave_sums = [alice_sum_d, bob_sum_d, carol_sum_d]\n",
    "dave_shapes = [alice_shape_d, bob_shape_d, carol_shape_d]\n",
    "\n",
    "# 此处不填 num_returns 等价于 num_returns=1，此时框架会按照 num_returns 的设定，把长度为 1 的 array-like 的返回值解包\n",
    "dave_avg = dave(calc_avg, num_returns=1)(dave_sums, dave_shapes)\n",
    "\n",
    "public_avg = sf.reveal(dave_avg)\n",
    "print(f'avg={sf.reveal(public_avg)}, type={type(public_avg)}')"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## 混合使用 SPU 密态设备编程"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "在这上节我们展示了 PYU 明文设备的用法和一些注意事项，并且展示了如何仅使用 PYU 做多方计算。但是仅仅使用 PYU 做多方计算也许不够安全，我们回顾一下公式：$$avg=\\frac{sum(a)+sum(b)+sum(c)}{n_a+n_b+n_c}$$\n",
    "\n",
    "在上述例子中，虽然每方都只向 dave 透露统计值，但是统计值仍然包含一些信息量：dave 知道了每一方的总和和数据量。这些信息是否“有害”就见仁见智了。\n",
    "\n",
    "那么如何让计算过程中的信息泄露最小化甚至完全消除呢？各类密文设备就是一种选择。\n",
    "\n",
    "\n",
    "我们仍然使用上面的例子，但是希望用 SPU 密态设备保护中间计算过程，消除中间信息泄露。\n",
    "\n",
    "与 PYU 不同，SPU 并不对应一个物理设备，而是由一组物理设备组成的一个虚拟的逻辑设备。SPU 的基础知识建议阅读[这篇文档](../developer/design/spu.rst)\n",
    "\n",
    "根据底层协议的不同，SPU 需要不同数量的物理设备组成，在这个例子中，我们选择需要 3 方参与的 ABY3 协议，并且选择 alice、bob、carol 三方组成 SPU。"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 26,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "\u001b[2m\u001b[36m(SPURuntime pid=3929327)\u001b[0m 2023-11-06 20:31:07.369 [info] [default_brpc_retry_policy.cc:DoRetry:52] socket error, sleep=1000000us and retry\n",
      "\u001b[2m\u001b[36m(SPURuntime pid=3929324)\u001b[0m 2023-11-06 20:31:07.438 [info] [default_brpc_retry_policy.cc:DoRetry:52] socket error, sleep=1000000us and retry\n",
      "\u001b[2m\u001b[36m(SPURuntime pid=3929327)\u001b[0m 2023-11-06 20:31:08.369 [info] [default_brpc_retry_policy.cc:LogHttpDetail:29] cntl ErrorCode '112', http status code '200', response header '', error msg '[E111]Fail to connect Socket{id=0 addr=127.0.0.1:54823} (0x0x3cbee40): Connection refused [R1][E112]Not connected to 127.0.0.1:54823 yet, server_id=0'\n",
      "\u001b[2m\u001b[36m(SPURuntime pid=3929327)\u001b[0m 2023-11-06 20:31:08.369 [info] [default_brpc_retry_policy.cc:DoRetry:75] aggressive retry, sleep=1000000us and retry\n",
      "\u001b[2m\u001b[36m(SPURuntime pid=3929324)\u001b[0m 2023-11-06 20:31:08.439 [info] [default_brpc_retry_policy.cc:LogHttpDetail:29] cntl ErrorCode '112', http status code '200', response header '', error msg '[E111]Fail to connect Socket{id=0 addr=127.0.0.1:58731} (0x0x36f6780): Connection refused [R1][E112]Not connected to 127.0.0.1:58731 yet, server_id=0'\n",
      "\u001b[2m\u001b[36m(SPURuntime pid=3929324)\u001b[0m 2023-11-06 20:31:08.439 [info] [default_brpc_retry_policy.cc:DoRetry:75] aggressive retry, sleep=1000000us and retry\n",
      "\u001b[2m\u001b[36m(SPURuntime pid=3929327)\u001b[0m 2023-11-06 20:31:09.369 [info] [default_brpc_retry_policy.cc:LogHttpDetail:29] cntl ErrorCode '112', http status code '200', response header '', error msg '[E111]Fail to connect Socket{id=0 addr=127.0.0.1:54823} (0x0x3cbee40): Connection refused [R1][E112]Not connected to 127.0.0.1:54823 yet, server_id=0 [R2][E112]Not connected to 127.0.0.1:54823 yet, server_id=0'\n",
      "\u001b[2m\u001b[36m(SPURuntime pid=3929327)\u001b[0m 2023-11-06 20:31:09.370 [info] [default_brpc_retry_policy.cc:DoRetry:75] aggressive retry, sleep=1000000us and retry\n",
      "\u001b[2m\u001b[36m(SPURuntime pid=3929324)\u001b[0m 2023-11-06 20:31:09.439 [info] [default_brpc_retry_policy.cc:LogHttpDetail:29] cntl ErrorCode '112', http status code '200', response header '', error msg '[E111]Fail to connect Socket{id=0 addr=127.0.0.1:58731} (0x0x36f6780): Connection refused [R1][E112]Not connected to 127.0.0.1:58731 yet, server_id=0 [R2][E112]Not connected to 127.0.0.1:58731 yet, server_id=0'\n",
      "\u001b[2m\u001b[36m(SPURuntime pid=3929324)\u001b[0m 2023-11-06 20:31:09.439 [info] [default_brpc_retry_policy.cc:DoRetry:75] aggressive retry, sleep=1000000us and retry\n",
      "\u001b[2m\u001b[36m(SPURuntime pid=3929325)\u001b[0m 2023-11-06 20:31:09.464 [info] [default_brpc_retry_policy.cc:DoRetry:69] not retry for reached rcp timeout, ErrorCode '1008', error msg '[E1008]Reached timeout=2000ms @127.0.0.1:54823'\n"
     ]
    }
   ],
   "source": [
    "import spu\n",
    "from secretflow.utils.testing import unused_tcp_port\n",
    "\n",
    "aby3_cluster_def = {\n",
    "    'nodes': [\n",
    "        {'party': 'alice', 'address': f'127.0.0.1:{unused_tcp_port()}'},\n",
    "        {'party': 'bob', 'address': f'127.0.0.1:{unused_tcp_port()}'},\n",
    "        {'party': 'carol', 'address': f'127.0.0.1:{unused_tcp_port()}'},\n",
    "    ],\n",
    "    'runtime_config': {\n",
    "        'protocol': spu.spu_pb2.ABY3,\n",
    "        'field': spu.spu_pb2.FM64,\n",
    "    },\n",
    "}\n",
    "\n",
    "my_spu = sf.SPU(aby3_cluster_def)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "我们仍然使用 PYU 加载本地数据："
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 27,
   "metadata": {},
   "outputs": [],
   "source": [
    "def agg(data):\n",
    "    return np.sum(data), np.shape(data)\n",
    "\n",
    "\n",
    "alice_sum, alice_shape = alice(agg, num_returns=2)(alice_data)\n",
    "bob_sum, bob_shape = bob(agg, num_returns=2)(bob_data)\n",
    "carol_sum, carol_shape = carol(agg, num_returns=2)(carol_data)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "目前 3 个 sum 和 3 个 shape 都分别在各自本地，因此是安全的。\n",
    "\n",
    "接下来需要把这些数据“放”到 `my_spu` 中计算均值，如果你了解一些[秘密分享](../developer/design/spu.rst#mpc-secure-multi-party-computation)协议就会知道，这一步骤需要各方本地做数据分片，并把部分分片发送到协议的其他方（也就是组成 SPU 的其他参与方）。\n",
    "\n",
    "这个步骤是相对复杂的，但是 SPU 设备在编程上屏蔽了这种复杂性：把数据从 PYU “放”到 SPU 就和在 PYU 之间传递一样简单。"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 28,
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "\u001b[2m\u001b[36m(_run pid=3909057)\u001b[0m INFO:jax._src.xla_bridge:Unable to initialize backend 'cuda': module 'jaxlib.xla_extension' has no attribute 'GpuAllocatorConfig'\n",
      "\u001b[2m\u001b[36m(_run pid=3909057)\u001b[0m INFO:jax._src.xla_bridge:Unable to initialize backend 'rocm': module 'jaxlib.xla_extension' has no attribute 'GpuAllocatorConfig'\n",
      "\u001b[2m\u001b[36m(_run pid=3909057)\u001b[0m INFO:jax._src.xla_bridge:Unable to initialize backend 'tpu': INVALID_ARGUMENT: TpuPlatform is not available.\n",
      "\u001b[2m\u001b[36m(_run pid=3909057)\u001b[0m INFO:jax._src.xla_bridge:Unable to initialize backend 'plugin': xla_extension has no attributes named get_plugin_device_client. Compile TensorFlow with //tensorflow/compiler/xla/python:enable_plugin_device set to true (defaults to false) to enable this.\n",
      "\u001b[2m\u001b[36m(_run pid=3909057)\u001b[0m WARNING:jax._src.xla_bridge:No GPU/TPU found, falling back to CPU. (Set TF_CPP_MIN_LOG_LEVEL=0 and rerun for more info.)\n",
      "\u001b[2m\u001b[36m(_run pid=3909321)\u001b[0m INFO:jax._src.xla_bridge:Unable to initialize backend 'cuda': module 'jaxlib.xla_extension' has no attribute 'GpuAllocatorConfig'\n",
      "\u001b[2m\u001b[36m(_run pid=3909321)\u001b[0m INFO:jax._src.xla_bridge:Unable to initialize backend 'rocm': module 'jaxlib.xla_extension' has no attribute 'GpuAllocatorConfig'\n",
      "\u001b[2m\u001b[36m(_run pid=3909321)\u001b[0m INFO:jax._src.xla_bridge:Unable to initialize backend 'tpu': INVALID_ARGUMENT: TpuPlatform is not available.\n",
      "\u001b[2m\u001b[36m(_run pid=3909321)\u001b[0m INFO:jax._src.xla_bridge:Unable to initialize backend 'plugin': xla_extension has no attributes named get_plugin_device_client. Compile TensorFlow with //tensorflow/compiler/xla/python:enable_plugin_device set to true (defaults to false) to enable this.\n",
      "\u001b[2m\u001b[36m(_run pid=3909321)\u001b[0m WARNING:jax._src.xla_bridge:No GPU/TPU found, falling back to CPU. (Set TF_CPP_MIN_LOG_LEVEL=0 and rerun for more info.)\n",
      "\u001b[2m\u001b[36m(_run pid=3908759)\u001b[0m INFO:jax._src.xla_bridge:Unable to initialize backend 'cuda': module 'jaxlib.xla_extension' has no attribute 'GpuAllocatorConfig'\n",
      "\u001b[2m\u001b[36m(_run pid=3908759)\u001b[0m INFO:jax._src.xla_bridge:Unable to initialize backend 'rocm': module 'jaxlib.xla_extension' has no attribute 'GpuAllocatorConfig'\n",
      "\u001b[2m\u001b[36m(_run pid=3908759)\u001b[0m INFO:jax._src.xla_bridge:Unable to initialize backend 'tpu': INVALID_ARGUMENT: TpuPlatform is not available.\n",
      "\u001b[2m\u001b[36m(_run pid=3908759)\u001b[0m INFO:jax._src.xla_bridge:Unable to initialize backend 'plugin': xla_extension has no attributes named get_plugin_device_client. Compile TensorFlow with //tensorflow/compiler/xla/python:enable_plugin_device set to true (defaults to false) to enable this.\n",
      "\u001b[2m\u001b[36m(_run pid=3908759)\u001b[0m WARNING:jax._src.xla_bridge:No GPU/TPU found, falling back to CPU. (Set TF_CPP_MIN_LOG_LEVEL=0 and rerun for more info.)\n"
     ]
    }
   ],
   "source": [
    "alice_sum_spu = alice_sum.to(my_spu)\n",
    "bob_sum_spu = bob_sum.to(my_spu)\n",
    "carol_sum_spu = carol_sum.to(my_spu)\n",
    "\n",
    "alice_shape_spu = alice_shape.to(my_spu)\n",
    "bob_shape_spu = bob_shape.to(my_spu)\n",
    "carol_shape_spu = carol_shape.to(my_spu)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "我们来看一下放到 SPU 以后，数据会变成什么样。\n",
    "\n",
    "跟 PYU 类似，把数据放到 SPU 以后，会得到一个 `secretflow.device.device.spu.SPUObject`, 这个 `SPUObject` 也是一个 SPU 中的数据的引用。和 PYU 不同，`SPUObject` 并不引用一个物理的数据，而是表示一个多方组成的逻辑数据，我们无法看到这个逻辑数据的内容，但是可以让这个逻辑数据在 `my_spu` 上执行计算。"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 29,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "(<secretflow.device.device.spu.SPUObject at 0x7f4bfef4c700>,\n",
       " <secretflow.device.device.spu.SPUObject at 0x7f4bfef07be0>)"
      ]
     },
     "execution_count": 29,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "alice_sum_spu, alice_shape_spu"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "同样来看一下 `SPUObject` 的内容。一个 `SPUObject` 中同样包含在所在的 spu 设备信息，此外还包含一个表示数据“形状”的 meta 信息和分片信息（不同组网模式下会有一些不同）。"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 30,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "{'device': <secretflow.device.device.spu.SPU object at 0x7f4bfefc9a00>,\n",
      " 'meta': ObjectRef(88543757a8df6d2fffffffffffffffffffffffff0100000001000000),\n",
      " 'shares_name': [ObjectRef(347cc60e0bb3da749fce953dc1bbfc1af4f50ef60100000001000000),\n",
      "                 ObjectRef(a02c24b8b7fc0a314bfe84a2ddef3fc302eba3980100000001000000),\n",
      "                 ObjectRef(a631fe8d231813bfab4501800d1eaed6ccbc32510100000001000000)]}\n"
     ]
    }
   ],
   "source": [
    "from pprint import pprint\n",
    "\n",
    "pprint(vars(alice_sum_spu))"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "现在所有数据都放到 `my_spu` 中，我们可以开始最终的均值计算了。整个计算都在密态进行，我们只将最终的结果公开出来。 "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 31,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "avg=13.374797821044922, type=<class 'numpy.ndarray'>\n"
     ]
    }
   ],
   "source": [
    "import jax.numpy as jnp\n",
    "\n",
    "\n",
    "def avg(sums, shapes):\n",
    "    sums = jnp.array(sums)\n",
    "    shapes = jnp.array(shapes)\n",
    "    res = jnp.sum(sums) / jnp.sum(shapes)\n",
    "    return res\n",
    "\n",
    "\n",
    "spu_sums = [alice_sum_spu, bob_sum_spu, carol_sum_spu]\n",
    "spu_shapes = [alice_shape_spu, bob_shape_spu, carol_shape_spu]\n",
    "\n",
    "spu_avg = my_spu(avg)(spu_sums, spu_shapes)\n",
    "\n",
    "public_avg = sf.reveal(spu_avg)\n",
    "print(f'avg={public_avg}, type={type(public_avg)}')"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "我们关注一下几个“细节”：\n",
    "1. SPU 执行的函数是使用 `jax.numpy` 写的，这个跟 SPU 的[设计](https://www.secretflow.org.cn/docs/spu/latest/en-US/getting_started/introduction)有关，目前只支持 JAX，支持的 JAX 算子的完整列表可以查看[文档](https://www.secretflow.org.cn/docs/spu/latest/en-US/reference/np_op_status)。\n",
    "\n",
    "2. 可以看到 SPU 的计算结果是正确的，但是精度不太高，这是因为我们用到了比较复杂的“除法”操作，可以考虑初始化 SPU 时更精细的设置 RuntimeConfig，\n",
    "RuntimeConfig 的完整选项可以查看 [RuntimeConfig 文档](https://www.secretflow.org.cn/docs/spu/latest/en-US/reference/runtime_config#runtimeconfig)\n",
    "\n",
    "\n",
    "接下来我们调整一下参数，再构造一个新的 spu 实例，看看能不能把计算精度提高："
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 32,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "\u001b[2m\u001b[36m(SPURuntime pid=3934768)\u001b[0m 2023-11-06 20:37:21.329 [info] [default_brpc_retry_policy.cc:DoRetry:52] socket error, sleep=1000000us and retry\n",
      "\u001b[2m\u001b[36m(SPURuntime pid=3934768)\u001b[0m 2023-11-06 20:37:22.329 [info] [default_brpc_retry_policy.cc:LogHttpDetail:29] cntl ErrorCode '112', http status code '200', response header '', error msg '[E111]Fail to connect Socket{id=0 addr=127.0.0.1:35155} (0x0x343db00): Connection refused [R1][E112]Not connected to 127.0.0.1:35155 yet, server_id=0'\n",
      "\u001b[2m\u001b[36m(SPURuntime pid=3934768)\u001b[0m 2023-11-06 20:37:22.330 [info] [default_brpc_retry_policy.cc:DoRetry:75] aggressive retry, sleep=1000000us and retry\n",
      "\u001b[2m\u001b[36m(SPURuntime pid=3934768)\u001b[0m 2023-11-06 20:37:23.330 [info] [default_brpc_retry_policy.cc:LogHttpDetail:29] cntl ErrorCode '112', http status code '200', response header '', error msg '[E111]Fail to connect Socket{id=0 addr=127.0.0.1:35155} (0x0x343db00): Connection refused [R1][E112]Not connected to 127.0.0.1:35155 yet, server_id=0 [R2][E112]Not connected to 127.0.0.1:35155 yet, server_id=0'\n",
      "\u001b[2m\u001b[36m(SPURuntime pid=3934768)\u001b[0m 2023-11-06 20:37:23.330 [info] [default_brpc_retry_policy.cc:DoRetry:75] aggressive retry, sleep=1000000us and retry\n",
      "\u001b[2m\u001b[36m(SPURuntime pid=3934769)\u001b[0m 2023-11-06 20:37:23.462 [info] [default_brpc_retry_policy.cc:DoRetry:69] not retry for reached rcp timeout, ErrorCode '1008', error msg '[E1008]Reached timeout=2000ms @127.0.0.1:46843'\n",
      "\u001b[2m\u001b[36m(SPURuntime pid=3934771)\u001b[0m 2023-11-06 20:37:23.484 [info] [default_brpc_retry_policy.cc:DoRetry:69] not retry for reached rcp timeout, ErrorCode '1008', error msg '[E1008]Reached timeout=2000ms @127.0.0.1:46843'\n"
     ]
    }
   ],
   "source": [
    "import secretflow as sf\n",
    "import spu\n",
    "from secretflow.utils.testing import unused_tcp_port\n",
    "\n",
    "aby3_cluster_def2 = {\n",
    "    'nodes': [\n",
    "        {'party': 'alice', 'address': f'127.0.0.1:{unused_tcp_port()}'},\n",
    "        {'party': 'bob', 'address': f'127.0.0.1:{unused_tcp_port()}'},\n",
    "        {'party': 'carol', 'address': f'127.0.0.1:{unused_tcp_port()}'},\n",
    "    ],\n",
    "    'runtime_config': {\n",
    "        'protocol': spu.spu_pb2.ABY3,\n",
    "        'field': spu.spu_pb2.FM64,\n",
    "        'fxp_fraction_bits': 24,\n",
    "        'fxp_div_goldschmidt_iters': 6,\n",
    "    },\n",
    "}\n",
    "\n",
    "my_spu2 = sf.SPU(aby3_cluster_def2)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 33,
   "metadata": {},
   "outputs": [],
   "source": [
    "# 把数据放到新的 spu 实例上\n",
    "alice_sum_spu = alice_sum.to(my_spu2)\n",
    "bob_sum_spu = bob_sum.to(my_spu2)\n",
    "carol_sum_spu = carol_sum.to(my_spu2)\n",
    "\n",
    "alice_shape_spu = alice_shape.to(my_spu2)\n",
    "bob_shape_spu = bob_shape.to(my_spu2)\n",
    "carol_shape_spu = carol_shape.to(my_spu2)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "在 `my_spu2` 上重新计算一次："
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 34,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "avg=13.375, type=<class 'numpy.ndarray'>\n"
     ]
    }
   ],
   "source": [
    "spu_sums = [alice_sum_spu, bob_sum_spu, carol_sum_spu]\n",
    "spu_shapes = [alice_shape_spu, bob_shape_spu, carol_shape_spu]\n",
    "\n",
    "spu_avg = my_spu2(avg)(spu_sums, spu_shapes)\n",
    "\n",
    "public_avg = sf.reveal(spu_avg)\n",
    "print(f'avg={public_avg}, type={type(public_avg)}')"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "当然除了通过 `reveal` 的方式公开结果，也可以把结果给到特定的“结果方”。SPU 到 PYU 的传递相当于把秘密分享分片集中到 PYU 中，重构出真实的数据，这个过程在编程中也非常简单："
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 35,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "<secretflow.device.device.pyu.PYUObject at 0x7f4bfef24880>"
      ]
     },
     "execution_count": 35,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "dave_avg = spu_avg.to(dave)\n",
    "\n",
    "dave_avg"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "到这里结果数据已经传递到了 dave，dave 可以进行本地计算或者保存等等操作。我们把 dave 的结果 `reveal` 出来看一下是否正确。"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 36,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "dave_avg=13.375\n"
     ]
    }
   ],
   "source": [
    "print(f'dave_avg={sf.reveal(dave_avg)}')"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### 算法封装"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "上面介绍了如何明密文混合编程，更多时候我们需要写一些封装好的算法或者功能模块，这里给出一些建议：\n",
    "1. 尽量不要在算法或功能模块内初始化集群和逻辑设备，把逻辑设备作为参数传入，可以让算法更通用；\n",
    "2. 有时候在算法中 PYU 明文设备和 SPU 密文设备是可替代的，对于这些部分建议写一套代码来支持（即用 jax 来写计算逻辑）。\n",
    "\n",
    "接下来我们把三方求均值做成一个可复用的功能模块。"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 37,
   "metadata": {},
   "outputs": [],
   "source": [
    "from typing import Dict\n",
    "from secretflow.device import Device\n",
    "\n",
    "\n",
    "def avg(files: Dict[Device, str], agg_device: Device, result_device: Device):\n",
    "    '''\n",
    "    参数：\n",
    "        files: PYU -> filename，表示每一方的文件路径\n",
    "        agg_device: 表示最终计算均值的虚拟设备，可以是 PYU 或者 SPU\n",
    "        result_device：表示获取结果的虚拟设备，一般来说是 PYU\n",
    "    '''\n",
    "\n",
    "    def _load_data(file_name):\n",
    "        return np.loadtxt(file_name, dtype=np.int32)\n",
    "\n",
    "    data = {device: device(_load_data)(fname) for device, fname in files.items()}\n",
    "\n",
    "    def _stat(data):\n",
    "        return np.sum(data), len(data)\n",
    "\n",
    "    sums = []\n",
    "    lengths = []\n",
    "    for device, datum in data.items():\n",
    "        sum, length = device(_stat, num_returns=2)(datum)\n",
    "        sums.append(sum.to(agg_device))\n",
    "        lengths.append(length.to(agg_device))\n",
    "\n",
    "    def _agg(sums, lengths):\n",
    "        sums = jnp.array(sums)\n",
    "        lengths = jnp.array(lengths)\n",
    "        return jnp.sum(sums) / jnp.sum(lengths)\n",
    "\n",
    "    result = agg_device(_agg)(sums, lengths)\n",
    "    return result.to(result_device)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "试着执行一下，先用明文设备 dave 作为聚合设备，并把最终结果传递给 alice"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 38,
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "INFO:jax._src.xla_bridge:Unable to initialize backend 'cuda': module 'jaxlib.xla_extension' has no attribute 'GpuAllocatorConfig'\n",
      "INFO:jax._src.xla_bridge:Unable to initialize backend 'rocm': module 'jaxlib.xla_extension' has no attribute 'GpuAllocatorConfig'\n",
      "INFO:jax._src.xla_bridge:Unable to initialize backend 'tpu': INVALID_ARGUMENT: TpuPlatform is not available.\n",
      "INFO:jax._src.xla_bridge:Unable to initialize backend 'plugin': xla_extension has no attributes named get_plugin_device_client. Compile TensorFlow with //tensorflow/compiler/xla/python:enable_plugin_device set to true (defaults to false) to enable this.\n",
      "WARNING:jax._src.xla_bridge:No GPU/TPU found, falling back to CPU. (Set TF_CPP_MIN_LOG_LEVEL=0 and rerun for more info.)\n"
     ]
    },
    {
     "data": {
      "text/plain": [
       "<secretflow.device.device.pyu.PYUObject at 0x7f4a9b8adc70>"
      ]
     },
     "execution_count": 38,
     "metadata": {},
     "output_type": "execute_result"
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "\u001b[2m\u001b[36m(_run pid=3909321)\u001b[0m result: 13.375\n"
     ]
    }
   ],
   "source": [
    "files = {\n",
    "    alice: './data/a.csv',\n",
    "    bob: './data/b.csv',\n",
    "    carol: './data/c.csv',\n",
    "}\n",
    "\n",
    "result = avg(files=files, agg_device=dave, result_device=alice)\n",
    "\n",
    "alice(print_data)('result', result)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "这次用 spu 作为聚合的逻辑设备，同样让结果在 alice 重建出来 "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 39,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "\u001b[2m\u001b[36m(_run pid=3909321)\u001b[0m my_spu: 13.374794006347656\n"
     ]
    },
    {
     "data": {
      "text/plain": [
       "<secretflow.device.device.pyu.PYUObject at 0x7f4bfef01af0>"
      ]
     },
     "execution_count": 39,
     "metadata": {},
     "output_type": "execute_result"
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "\u001b[2m\u001b[36m(_run pid=3909321)\u001b[0m my_spu2: 13.375\n"
     ]
    }
   ],
   "source": [
    "# 用低精度的 SPU 做聚合\n",
    "result = avg(files=files, agg_device=my_spu, result_device=alice)\n",
    "alice(print_data)('my_spu', result)\n",
    "\n",
    "# 用高精度的 SPU 做聚合\n",
    "result = avg(files=files, agg_device=my_spu2, result_device=alice)\n",
    "alice(print_data)('my_spu2', result)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## 更多密态设备"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### HEU 同态加密设备"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "HEU 是隐语的另一个密态设备。HEU 的使用方式与 SPU 或 PYU 差异较大，HEU 目前不支持可编程能力，即：不能定义一个函数然后交给 HEU 执行。HEU 目前更像是一个同态加密库，SecretFlow 对他做了一些封装，使得它可以和其他逻辑设备联合使用。\n",
    "\n",
    "\n",
    "HEU 的初始化也相对复杂，需要对同态加密有一定的了解，建议先阅读 [HEU 文档](https://www.secretflow.org.cn/docs/heu/latest/zh-Hans/getting_started/quick_start)了解基础概念。在编程中需要特别注意 `sk_keeper` 的设置，他拥有私钥可以解密数据，因此大多数时候加密数据不应该交给 `sk_keeper` 来计算。"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 40,
   "metadata": {},
   "outputs": [],
   "source": [
    "import spu\n",
    "\n",
    "my_heu = sf.HEU(\n",
    "    {\n",
    "        'sk_keeper': {'party': 'alice'},\n",
    "        'evaluators': [{'party': 'bob'}, {'party': 'carol'}, {'party': 'dave'}],\n",
    "        'mode': 'PHEU',\n",
    "        'he_parameters': {\n",
    "            'schema': 'paillier',\n",
    "            'key_pair': {'generate': {'bit_size': 2048}},\n",
    "        },\n",
    "    },\n",
    "    spu.spu_pb2.FM128,\n",
    ")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "数据转移到 HEU 时需要一个 `HEUMoveConfig` ，它包含两个参数：\n",
    "1. heu_encoder： 用来指定数据编码方式，根据原始数据的类型选择用整型编码器或者浮点数编码器；\n",
    "2. heu_dest_party：表示数据实际传输到哪一方，HEU 会自动决定是否需要加密：如果 `heu_dest_party` 与 原始数据所在的 `PYU` 一致就**不会加密**数据，如果数据出现了跨域传输 HEU 就会把它**加密**。\n",
    "\n",
    "在编程过程中，有时候需要注意 `HEUObject` 是明文还是密文状态，对于加法同态来说，乘法只能是一个密文乘以一个明文（可以理解为多个密文的和）。"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 41,
   "metadata": {},
   "outputs": [],
   "source": [
    "from secretflow.device.device.heu import HEUMoveConfig\n",
    "from heu import phe\n",
    "\n",
    "move_config = HEUMoveConfig(\n",
    "    heu_dest_party='dave',\n",
    "    heu_encoder=phe.IntegerEncoder(my_heu.schema),\n",
    ")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "定义好 `move_config` 后，我们就可以用类似其他逻辑设备的方式把数据放到 `HEU` 了："
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 42,
   "metadata": {},
   "outputs": [],
   "source": [
    "alice_sum_heu = alice_sum.to(my_heu, move_config)\n",
    "bob_sum_heu = bob_sum.to(my_heu, move_config)\n",
    "carol_sum_heu = carol_sum.to(my_heu, move_config)\n",
    "\n",
    "alice_shape_heu = alice_shape.to(my_heu, move_config)\n",
    "bob_shape_heu = bob_shape.to(my_heu, move_config)\n",
    "carol_shape_heu = carol_shape.to(my_heu, move_config)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "同样看一下 `HEU` 里面的数据是什么样的，因为我们的数据进行了跨设备移动，所以得到的 HEUObject 是加密状态的："
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 43,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "{'data': ObjectRef(f0133118a11aef4156f60ab8810d57c91a9063380100000001000000),\n",
      " 'device': <secretflow.device.device.heu.HEU object at 0x7f4bfef18b20>,\n",
      " 'is_plain': False,\n",
      " 'location': 'dave'}\n"
     ]
    }
   ],
   "source": [
    "pprint(vars(alice_sum_heu))"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "与其他逻辑设备不同，可以直接在 `HEUObject` 执行运算："
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 44,
   "metadata": {},
   "outputs": [],
   "source": [
    "total_sum_heu = alice_sum_heu + bob_sum_heu + carol_sum_heu\n",
    "total_shape_heu = alice_shape_heu + bob_shape_heu + carol_shape_heu"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "计算的结果仍然是一个 `HEUObject`："
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 45,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "{'data': ObjectRef(7ef10a5289b2683b2b83696828bdba01a38857140100000001000000),\n",
      " 'device': <secretflow.device.device.heu.HEU object at 0x7f4bfef18b20>,\n",
      " 'is_plain': False,\n",
      " 'location': 'dave'}\n"
     ]
    }
   ],
   "source": [
    "pprint(vars(total_sum_heu))"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "HEU 目前不支持除法，所以需要把除法操作交给其他逻辑设备。\n",
    "\n",
    "数据在 `HEU` 到其他设备的传递也使用同样的 API，但是根据目标设备的不同(`PYU` 或者 `SPU`)，背后的实现方式非常不同。在编程的过程中你只需要知道：数据在 `SPU` 和 `HEU` 之间传递是安全的，中间没有任何数据信息泄露。\n",
    "\n",
    "下面的例子中，我们把数据放到 PYU 设备上执行除法操作："
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 46,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "(<secretflow.device.device.pyu.PYUObject at 0x7f4bfefc9310>,\n",
       " <secretflow.device.device.pyu.PYUObject at 0x7f4bfefc9220>)"
      ]
     },
     "execution_count": 46,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "total_sum_alice = total_sum_heu.to(alice)\n",
    "total_shape_alice = total_shape_heu.to(alice)\n",
    "\n",
    "total_sum_alice, total_shape_alice"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 47,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "avg_alice=[13.375]\n"
     ]
    }
   ],
   "source": [
    "def calc_avg(total_sum, total_shape):\n",
    "    return total_sum / total_shape\n",
    "\n",
    "\n",
    "avg_alice = alice(calc_avg)(total_sum_alice, total_shape_alice)\n",
    "\n",
    "print(f\"avg_alice={sf.reveal(avg_alice)}\")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "HEUObject 支持加法、减法、乘法和矩阵乘法（+, -, *, @）几种运算，加减法的两个操作数都可以是明文或密文，但是乘法的两个操作数必须有一个是明文。因为包括我们上面选用的 paillier 在内，HEU 目前只支持同态加法，详细算法列表可以查看 [HEU 文档](https://www.secretflow.org.cn/docs/heu/latest/zh-Hans/getting_started/algo_choice)。"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### TEEU 可信执行环境逻辑设备"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "可信执行环境(Trusted Execution Environment，TEE) 是一种基于硬件的隐私保护技术。它保证了执行代码的真实性，运行时状态(如 CPU 寄存器、内存和敏感 I/O)的完整性，以及存储在内存中的代码、数据和运行时状态的机密性。此外，还应能够向第三方提供远程认证，以证明其可靠性。\n",
    "\n",
    "TEEU(TEE processing Unit)是 SecretFlow 中的 TEE 逻辑设备，通过 TEEU，用户可以方便的把数据放在TEE内进行计算，并且达到保护数据完整和安全的目的。\n",
    "\n",
    "TEEU 的编码体验与 PYU 类似，只要注意 `to` 的 `allow_funcs` 参数，表示允许这个数据在 `TEEU` 中执行的逻辑。但是由于 TEE 的特点，TEEU 的部署相对复杂，建议仔细阅读[这篇文档](https://www.secretflow.org.cn/docs/secretflow/latest/zh-Hans/tutorial/teeu)。\n",
    "\n",
    "示例代码：\n",
    "```python\n",
    "from secretflow.device import TEEU\n",
    "\n",
    "def average(data):\n",
    "    return np.average(data, axis=1)\n",
    "\n",
    "# mrenclave can be omitted in simulation mode.\n",
    "teeu = TEEU('carol', mr_enclave='')\n",
    "\n",
    "# Transfer data to teeu.\n",
    "a_teeu = alice_data.to(teeu, allow_funcs=average)\n",
    "b_teeu = bob_data.to(teeu, allow_funcs=average)\n",
    "\n",
    "# TEEU runs average.\n",
    "avg_val = teeu(average)([a_teeu, b_teeu])\n",
    "```"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3 (ipykernel)",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.8.17"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 4
}
